"""
# Reference

[Densely Connected Convolutional Networks]()

This model is based on the following implementations:
- https://github.com/calmisential/DenseNet_TensorFlow2/blob/master/densenet.py

"""
from tensorflow.keras import Model, layers, activations, Sequential


class ConvBlock(layers.Layer):
    def __init__(self, growth_rate=32):
        super(ConvBlock, self).__init__()
        self.block_bn1 = layers.BatchNormalization()
        self.block_relu1 = layers.ReLU()
        self.block_conv1 = layers.Conv2D(4 * growth_rate, 1, 1, "same")

        self.block_bn2 = layers.BatchNormalization()
        self.block_relu2 = layers.ReLU()
        self.block_conv2 = layers.Conv2D(growth_rate, 3, 1, "same")

        self.block_concat = layers.Concatenate(axis=3)

    def call(self, inputs, training=None, **kwargs):
        x = self.block_bn1(inputs, training=training)
        x = self.block_relu1(x)
        x = self.block_conv1(x)

        x = self.block_bn2(x, training=training)
        x = self.block_relu2(x)
        x = self.block_conv2(x)

        x = self.block_concat([inputs, x])
        return x


class DenseBlock(layers.Layer):
    def __init__(self, num_layers, growth_rate):
        super(DenseBlock, self).__init__()
        self.num_layers = num_layers
        self.growth_rate = growth_rate
        self.dense_layer = self._make_layer()

    def _make_layer(self):
        features_list = []
        for i in range(self.num_layers):
            features_list.append(ConvBlock(self.growth_rate))
        return Sequential(features_list)

    def call(self, inputs, training=None, **kwargs):
        x = self.dense_layer(inputs)
        return x


class TransitionLayer(layers.Layer):
    def __init__(self, out_channels):
        super(TransitionLayer, self).__init__()
        self.trans_bn1 = layers.BatchNormalization()
        self.trans_relu1 = layers.ReLU()
        self.trans_conv1 = layers.Conv2D(out_channels, 1, 1, "same")
        self.trans_pool1 = layers.AveragePooling2D(2, 2, "same")

    def call(self, inputs, training=None, **kwargs):
        x = self.trans_bn1(inputs)
        x = self.trans_relu1(x)
        x = self.trans_conv1(x)
        x = self.trans_pool1(x)
        return x


class DenseNet(Model):
    def __init__(self, num_init_features, growth_rate, block_layers, compression_rate, drop_rate, num_class=200):
        super(DenseNet, self).__init__()
        self.conv1 = layers.Conv2D(num_init_features, 7, 2, "same")
        self.bn1 = layers.BatchNormalization()
        self.relu1 = layers.ReLU()
        self.pool1 = layers.MaxPool2D(3, 2, "same")

        self.num_channels = num_init_features
        self.dense_block_1 = DenseBlock(num_layers=block_layers[0], growth_rate=growth_rate)
        self.num_channels += growth_rate * block_layers[0]
        self.num_channels = compression_rate * self.num_channels
        self.transition_1 = TransitionLayer(out_channels=int(self.num_channels))

        self.dense_block_2 = DenseBlock(num_layers=block_layers[1], growth_rate=growth_rate)
        self.num_channels += growth_rate * block_layers[1]
        self.num_channels = compression_rate * self.num_channels
        self.transition_2 = TransitionLayer(out_channels=int(self.num_channels))

        self.dense_block_3 = DenseBlock(num_layers=block_layers[2], growth_rate=growth_rate)
        self.num_channels += growth_rate * block_layers[2]
        self.num_channels = compression_rate * self.num_channels
        self.transition_3 = TransitionLayer(out_channels=int(self.num_channels))

        # self.dense_block_4 = DenseBlock(num_layers=block_layers[3], growth_rate=growth_rate)

        self.bn2 = layers.BatchNormalization()
        self.relu2 = layers.ReLU()
        self.avg_pool = layers.GlobalAveragePooling2D()
        self.fc = layers.Dense(units=num_class, activation=activations.softmax)

    def call(self, inputs, training=None, **kwargs):
        x = self.conv1(inputs)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.pool1(x)

        x = self.dense_block_1(x, training=training)
        x = self.transition_1(x, training=training)
        x = self.dense_block_2(x, training=training)
        x = self.transition_2(x, training=training)
        x = self.dense_block_3(x, training=training)
        x = self.transition_3(x, training=training)
        # x = self.dense_block_4(x, training=training)

        x = self.bn2(x)
        x = self.relu2(x)
        x = self.avg_pool(x)
        x = self.fc(x)

        return x


def densenet121():
    # return DenseNet(64, 32, [6, 12, 24, 16], 0.5, 0.5)
    return DenseNet(64, 32, [3, 6, 12, 8], 0.5, 0.5)


def densenet169():
    return DenseNet(64, 32, [6, 12, 32, 32], 0.5, 0.5)


def densenet201():
    return DenseNet(64, 32, [6, 12, 48, 32], 0.5, 0.5)


def densenet264():
    return DenseNet(64, 32, [6, 12, 64, 48], 0.5, 0.5)
